import numpy as np
import torch
import torch.nn.functional as F
import torch.optim as optim

from defense.base_defense import BaseDefense

class Defense(BaseDefense):
    def __init__(self, model_name, config):
        super().__init__('NAT', model_name, config, retrain = True)

    #进行LLC攻击
    def LLC_attack(self, images):
        epsilon = np.random.normal(loc = self.config['eps_mu'], scale = self.config['eps_sigma'], size = [images.size(0)]) / 255.0
        epsilon = np.clip(np.abs(epsilon), self.config['clip_min'], self.config['clip_max'])
        var_epsilon = torch.from_numpy(epsilon).float().to(self.model.device)

        images.requires_grad = True

        self.model.eval()
        output = self.model(images)
        llcs = torch.argmin(output, dim = 1)
        loss = F.cross_entropy(output, llcs)
        grad = torch.autograd.grad(loss, images)[0]

        images.requires_grad = False

        grad_sign = torch.sign(grad)

        with torch.no_grad():
            advs = []
            for i in range(images.size(0)):
                adv = images[i] - var_epsilon[i] * grad_sign[i]
                adv = torch.clamp(adv, min = 0.0, max = 1.0)
                advs.append(adv)
            ret = torch.stack(advs)
        ret = torch.clamp(ret, min = 0.0, max = 1.0)

        return ret

    #单轮防御训练
    def train(self, loader, opt, epoch):
        for ind, (images, labels) in enumerate(loader):
            images = images.to(self.model.device)
            labels = labels.to(self.model.device)

            #进行LLC攻击生成混淆图像
            advs = self.LLC_attack(images)

            self.model.train()

            output = self.model(images)
            loss = F.cross_entropy(output, labels) #原始图像损失

            adv_output = self.model(advs)
            adv_loss = F.cross_entropy(adv_output, labels) #混淆图像损失

            total_loss = (loss + self.config['adv_ratio'] * adv_loss) / (1.0 + self.config['adv_ratio']) #总损失

            opt.zero_grad()
            total_loss.backward()
            opt.step()

            ind += 1
            print('\rEpoch {}: [batch: {}/{} ({:.0f}%)] \tLoss: {:.4f} AdvLoss: {:.4f} TotalLoss: {:.4f}'.format(
                epoch, ind, len(loader),
                ind / len(loader) * 100, loss, adv_loss, total_loss), end=' ')

        print()

    def defend(self):
        batch_size = self.model.training_params['batch_size']
        #加载训练数据
        train_loader, valid_loader = self.model.get_train_valid_loader(batch_size = batch_size,
                                                                        valid_size = 0.1,
                                                                        transform_train = True,
                                                                        shuffle = True)
        #获取optimizer
        opt = self.model.get_optimizer()

        best_ratio = 0
        best_state_dict = None
        for epoch in range(self.model.training_params['num_epochs']):
            self.train(train_loader, opt, epoch)
            ratio = self.model.test(valid_loader)
            self.model.adjust_optimizer(opt, epoch)
            if ratio > best_ratio:
                best_ratio = ratio
                best_state_dict = self.model.state_dict_clone()
                print('Validate: Accuracy increased')
            else:
                print('Validate: Accuracy decreased')
            print()
        return best_state_dict
